#-*- coding:utf-8 -*-

import os,cv2,time,shutil,random


class VideoLabel():
    def __init(self):
        self.src = ""
        self.dst = ""

    def convert(self,size, box):
        # rx:  归一化的中心点x坐标计算公式：(x2 + x1) / 2.0 / iw
        # ry:  归一化的中心点y坐标计算公式：(y2 + y1) / 2.0 / ih
        # rw/2:归一化的目标框宽度的计算公式：(x2 - x1)/ 2.0 / iw
        # rh/2:归一化的目标框高度的计算公式：(y2 - y1)/ 2.0 / ih

        # x2 = (rx+rw/2.0)*iw
        # x1 = (rx-rw/2.0)*iw
        # y2 = (ry+rh/2.0)*ih
        # y1 = (ry1-h/2.0)*ih

        dw = 1. / size[0]
        dh = 1. / size[1]

        x = (box[0] + box[1]) / 2.0
        y = (box[2] + box[3]) / 2.0
        w = box[1] - box[0]
        h = box[3] - box[2]
        x = x * dw
        w = w * dw
        y = y * dh
        h = h * dh
        return x, y, w, h

    def convertR(self,rx, ry, rw, rh, iw, ih):
        # rx:  归一化的中心点x坐标计算公式：(x2 + x1) / 2.0 / iw
        # ry:  归一化的中心点y坐标计算公式：(y2 + y1) / 2.0 / ih
        # rw/2:归一化的目标框宽度的计算公式：(x2 - x1)/ 2.0 / iw
        # rh/2:归一化的目标框高度的计算公式：(y2 - y1)/ 2.0 / ih
        rx = rx * iw
        ry = ry * ih
        rw = rw * iw
        rh = rh * ih

        x2 = (rx + rw / 2.0)
        x1 = (rx - rw / 2.0)
        y2 = (ry + rh / 2.0)
        y1 = (ry - rh / 2.0)

        # x2 = (rx + rw / 2.0) * iw
        # x1 = (rx - rw / 2.0) * iw
        # y2 = (ry + rh / 2.0) * ih
        # y1 = (ry - rh / 2.0) * ih
        # print(x1,x2,y1,y2)
        x, y, w, h = x1, y1, x2 - x1, y2 - y1
        return x, y, w, h

    def removeSmallImage(self,src,dst):
        for file in os.listdir(src):
            if(file.endswith(".txt") and (not file.startswith("cl"))):
                img = cv2.imread(os.path.join(src, file.split(".")[0]+".jpg"))
                _, iw, ih = img.shape[::-1]
                locs = []
                with open(os.path.join(src,file),"r") as f:
                    lines = f.readlines()
                    for line in lines:
                        label,x,y,w,h = map(float,line.rstrip().split(" "))
                        rx,ry,rw,rh = self.convertR(x,y,w,h,iw,ih)
                        if(int(label==0)):
                            if (rw > 60 and rh > 100):
                                locs.append([label,x,y,w,h])
                        else:
                            locs.append([label, x, y, w, h])

                name = file.split(".")[0]
                with open(os.path.join(dst,name+".txt"),'w') as f:
                    for loc in locs:
                        label, x, y, w, h = loc
                        x, y, w, h = map("{:.6}".format,[x, y, w, h])
                        loc = [str(int(label)),x, y, w, h]
                        loc = " ".join(map(str,loc))+"\n"
                        print(loc)
                        f.write(loc)



    def videoMsg(self,cap):
        fps = cap.get(cv2.CAP_PROP_FPS)
        fourcc = cv2.VideoWriter_fourcc(*"MJPG")
        w, h = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        return fps,fourcc,[w,h]

    def randInt(self,start,end):
        return random.randint(int(start),int(end))

    def makeYoloTrainDataFormat(self,src, dst):
        timeToday = str(time.localtime().tm_mon) + "_" + str(time.localtime().tm_mday)
        # timeToday="9_23"
        trainName = src.split("/")[-1] + "Train" + timeToday + ".txt"

        trainTxt = open(os.path.join(src, trainName), "w")
        ii = 0
        for num, i in enumerate(os.listdir(src)):
            if ((not i.endswith("g")) or (i.endswith("G"))):
                continue
            imgName = i
            imgTxt = imgName.split(".")[0] + ".txt"
            if (not os.path.isfile(os.path.join(os.path.abspath(src), i))):
                print(imgTxt)
                ii += 1

            txtName = i.split(".")[0] + ".txt"
            trainTxt.write(os.path.join(os.path.abspath(src), i) + "\n")
        trainTxt.close()
        print(ii)

    def nms(self,loc):

        if(len(loc)==1):
            return True

        for num,l in enumerate(loc):
            x,y,w,h = l

    def checkSame(self,src):
        txtFile ,jpgFile = [],[]
        for file in os.listdir(src):
            if(file.endswith(".txt")):
                if(not file.startswith("cl")):
                    if(not file.startswith("tr")):
                        txtFile.append(os.path.join(src,file.split(".")[0]))

            if(file.endswith("g") or file.endswith("G")):
                jpgFile.append(os.path.join(src,file.split(".")[0]))

        needRemove = []

        for file in txtFile:
            jpgName = file.split(".")[0]
            if(jpgName not in jpgFile):
                needRemove.append(jpgName)

        for file in jpgFile:
            txtName = file.split(".")[0]
            if(txtName not in txtFile):
                needRemove.append(txtName)
        # print(needRemove)
        for file in needRemove:
            print(file,"removed")
            os.system("rm "+file+"*")


    def checkOutLabels(self,src):
        needRemove = []
        return

        imgList ,txtList = [],[]
        for file in os.listdir(src):

            if(file.endswith("jpg")):
                imgList.append(file)
            # if(file.endswith(".txt") and (not file.startswith("cl")) and ("lab" not in file)):
            #     txtList.append(file)

        for file in imgList:
            txtName = file.replace("jpg","txt")
            if(not os.path.isfile(os.path.join(src,txtName))  or not  (os.path.isfile(os.path.join(src,file))) ):
                try:
                    needRemove.append(os.path.join(src,txtName))
                    needRemove.append(os.path.join(src, file))
                    os.remove(os.path.join(src,txtName))
                    os.remove(os.path.join(src, file))
                except:
                    print (file," need remove")

        for file in os.listdir(src):
            if(file.endswith(".jpg")):
                fileTxt = file.replace(".jpg",".txt")
                fileImg = file

                try:
                    with open(os.path.join(src,fileTxt),"r") as f:
                        lines = f.readlines()
                        if (len(lines)==0):
                            needRemove.append(os.path.join(src,fileTxt))
                except:
                    if(os.path.isfile(os.path.join(src,fileTxt))):
                        print(fileTxt, "removed")
                        os.remove(os.path.join(src,fileTxt))
                    if(os.path.isfile(os.path.join(src,fileTxt.replace(".txt",".jpg")))):
                        os.remove(os.path.join(src,fileTxt.replace(".txt",".jpg")))

        for file in needRemove:
            print(file)
            if(os.path.isfile(file)):
                os.remove(file)
            if(os.path.isfile(file.replace(".txt",".jpg"))):
                os.remove(file.replace(".txt",".jpg"))

        txtList ,jpgList,needRemove = [],[],[]
        for file  in os.listdir(src):
            if (file.endswith("txt")):
                if(not os.path.isfile(file)):
                    needRemove.append(file)
            if (file.endswith("jpg")):
                if (not os.path.isfile(file)):
                    needRemove.append(file)
            # if(file.endswith("txt")):
            #     txtList.append(file)
            # if(file.endswith("jpg")):
            #     jpgList.append(file)

        # for txt in txtList:
        #     jpg = txt.replace("txt","jpg")
        #     if(os.path.isfile(os.path.join(src,jpg))):


        for file in needRemove:
            file = os.path.join(src,file)
            os.remove(file)
            print(file,"removed")

    def removeNullTxt(self,labelImgPath):
        for file in os.listdir(labelImgPath):
            if(file.endswith(".txt")):
                filePath = os.path.join(labelImgPath,file)
                with open(filePath,'r') as f:
                    lines = f.readlines()
                    if(len(lines) == 0):
                        os.remove(os.remove(filePath))
                        os.remove(os.remove(filePath.replace(".txt",".jpt")))
                        print(filePath ,"has removed")



    def checkTrainValLabels(self,trainLabel):
        labels,needRemove = [],[]
        with open(trainLabel,"r") as f:
            lines = f.readlines()
            for line in lines:
                line = line.rstrip()
                if(not os.path.isfile(line)):
                    print (line)
                    needRemove.append(line)

                    txtFile = line.replace("jpg","txt")
                    needRemove.append(txtFile)
                if(not os.path.isfile(line.replace("jpg","txt"))):
                    print (line,"2")
                    needRemove.append(line)
                    jpgFile = line.replace("txt","jpg")
                    needRemove.append(jpgFile)

        for file in needRemove:
            os.replace(file)
            os.replace(file)
            os.replace(file.replace("jpg",'txt'))
            os.replace(file)

            print (file)

        print (set(needRemove))

        with open(trainLabel.replace(".txt","_new.txt"),'w') as f:
            with open(trainLabel,"r") as trainFile:
                lines = trainFile.readlines()
                for line in lines:
                    line = line.rstrip()
                    if ( not (line in needRemove) ):
                        f.write(line+"\n")

    def gauBlur(self,img,kernelSize=(3,3),sigma=1.5):
        return cv2.GaussianBlur(img,kernelSize,sigma)

    def getMax(self,x,distance,target=0,model=1):
        if (model == 1):
            return max(x-distance,target)
        else:
            return max(x - distance, target)


    def getMin(self,x,distance,target):
        return min(x+distance,target)

    def paddingLoc(self,x,y,w,h,iw,ih,distance=30):
        return self.getMax(x,distance),self.getMax(y,distance),self.getMin(w,2*distance,iw),self.getMin(h,2*distance,ih)

    def gauAugmentation(self,src,dst=''):
        imgList = []

        for file in os.listdir(src):
            if(file.endswith(".jpg")):
                txtName = file.replace(".jpg",".txt")
                if( (not os.path.isfile(os.path.join(src,file))) or (not os.path.isfile(  os.path.join(src,txtName) ))):
                    try:
                        os.remove(os.path.join(src,file))
                        os.remove(os.path.join(src,txtName))
                        print (file +" has removed")
                        # continue
                    except:
                        continue

        padding = 50
        for file in os.listdir(src):
            if(file.endswith(".jpg")):
                img = cv2.imread(os.path.join(src,file))
                _,iw,ih = img.shape[::-1]
                txtName = file.replace(".jpg",".txt")

                with open(os.path.join(src,txtName),"r") as f:
                    lines = f.readlines()
                    for line in lines:
                        line = line.rstrip()
                        label,rx,ry,rw,rh = list(map(float,line.split(" ")))
                        x, y, w, h = self.convertR(rx,ry,rw,rh,iw,ih)
                        x, y, w, h = list(map(int,[x, y, w, h]))
                        print (x,y,w,h)
                        addX,addY,addW,addH = self.paddingLoc(x,y,w,h,iw,ih,padding)
                        roi = img[slice(addY,addY+addH),slice(addX,addX+addW)]
                        roiGau = self.gauBlur(roi,(5,5))
                        img[slice(addY,addY+addH),slice(addX,addX+addW)] = roiGau
                        # cv2.rectangle(img,(x,y),(x+w,y+h),(255,0,0),2,2)
                        cv2.rectangle(img,(addX,addY),(addX+addW,addY+addH),(255,0,0),2,2)

            cv2.imshow("img", img)
            cv2.waitKey(3000)


    def mkClassData(self,src,dst,data):

        for file in os.listdir(src):
            if(file.endswith(".jpg")):
                shutil.copy(os.path.join(src,file),os.path.join(dst,"adv_"+str(data)+"_"+file))

    # def mkClassDataList(self,src):
    #     for file in os.listdir(src):

    def readFPS(self,videoPath,dstPath):
        cap = cv2.VideoCapture(videoPath)
        # fps, fourcc, [w, h] = self.videoMsg(cap)
        ret,frame = cap.read()
        videoName = videoPath.split("/")[-1].split(".")[0]
        fpsNumber = 0
        # distanceTime = int(fps)#每间隔1s保存一次图片
        distanceTime =25#每间隔1s保存一次图片
        while(ret):
            try:
                ret,frame = cap.read()
            except:
                print("error")
                break
            if(frame is not None):
                # fpsStop = int(fpsNumber/distanceTime)
                # if(fpsStop == 11):
                #     print("stop")
                if(fpsNumber % distanceTime == 0):
                    cv2.imwrite(os.path.join(dstPath,videoName+"fps"+str(int(fpsNumber/distanceTime)).zfill(6)+".jpg"),frame)
                    print(int(fpsNumber/distanceTime))
                    try:
                        tmp = cv2.resize(frame,(960,540))
                        cv2.imshow("frame",tmp)
                    except:
                        continue
                    cv2.waitKey(1)

            fpsNumber+=1


    def mkpath(self,path):
        if(os.path.exists(path)):
            return path
        os.mkdir(path)
        return path

    def getRoi(self,src,dst):
        ends = ["jpg","jpeg","png"]
        for file in os.listdir(src):
            if(file.startswith("train") or file.startswith("val")):
                continue
            if(file.endswith("g")):
                continue
            if(file.startswith("cl")):
                continue

            # if(not os.path.exists(os.path.join(src,file))):
            #     continue
            # if (not os.path.exists(os.path.join(src, file.replace(".txt",".jpg")))):
            #     continue

            txtPath = os.path.join(src,file)
            for end in ends:
                imgPath = os.path.join(src,file.replace("txt",end))
                if(os.path.exists(os.path.join(src,imgPath))):
                    break

            with open(txtPath, 'r') as f:
                # try:
                img = cv2.imread(imgPath)
                # except:
                #     continue
                # print(txtPath,imgPath)
                imgTmp = img.copy()

                _, iw, ih = img.shape[::-1]
                # try:
                lines = f.readlines()
                # except:
                #     print("here")

                for num, line in enumerate(lines):
                    label, xr, yr, wr, hr = list(map(float, line.rstrip().split(" ")))
                    # label = '1'
                    x, y, w, h = self.convertR(xr, yr, wr, hr, iw, ih)
                    x, y, w, h = list(map(int, [x, y, w, h]))
                    if(w<100 or h<100):
                        continue

                    print(w,h)
                    # if(int(label) == 0):
                    if (True):
                        cv2.rectangle(imgTmp, (x, y), (x + w, y + h), (0, 0, 255), 4, 4)
                        roi = img[slice(y, y + h), slice(x, x + w)]
                        # labelPath = os.path.join(os.path.join(dst, labels[int(label)]),i.split(".")[0] + "_" + str(num).zfill(6) + ".jpg")
                        labelPath = os.path.join(self.mkpath(os.path.join(dst,str(int(label)))),file.split(".")[0]+str(num).zfill(2)+".jpg")

                        cv2.imwrite(labelPath, roi)
                        cv2.imshow("frame",imgTmp)
                        cv2.waitKey(1)


    def removeImg(self,src,dst):
        for file in os.listdir(src):
            if(file.endswith("g") or file.endswith("G")):
                txtPath = os.path.join(src,file.split(".")[0]+".txt")
                if( not os.path.isfile(txtPath)):
                    shutil.move(os.path.join(src,file),os.path.join(dst,file))

videoLabel = VideoLabel()



